# service/parse/s3_upload.py
import io
import os

import boto3
from botocore.exceptions import BotoCoreError, ClientError

from common.exception.exception import BusiException
from common.exception.exception_enum import BusiExceptionEnum
from common.log.logging import log
from common.oss.oss import OssApi, OssObject
from common.utils.id_utils import get_uuid


class OssS3(OssApi):
    def __init__(self, access_key: str, secret_key: str, endpoint: str, default_bucket: str):
        self.s3_client = boto3.client(
            's3',
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            endpoint_url=endpoint
        )
        self.default_bucket = default_bucket

    def upload(self, filename: str, stream: io.IOBase) -> OssObject:
        """
        文件上传到默认的bucket
        :param filename: 文件名称
        :param stream: 上传的文件
        :return: OssObject
        """
        return self.upload_to_bucket(self.default_bucket, filename, stream)

    def upload_to_bucket(self, bucket_name: str, filename: str, stream: io.IOBase) -> OssObject:
        """
        上传文件到 AWS S3。

        参数:
        - file: io.IOBase 对象，表示上传的文件。

        返回:
        - OssObject 对象，包含桶名称和对象 ID。
        """
        if not stream:
            raise BusiException(BusiExceptionEnum.NO_FILE_PART)

        try:
            if not self.bucket_exists(bucket_name):
                self.create_bucket(bucket_name)
            unique_filename = f"{get_uuid()}{os.path.splitext(filename)[1]}"
            self.s3_client.upload_fileobj(stream, bucket_name, unique_filename)
            log.info(f'File {filename} successfully uploaded to {bucket_name}')
            return OssObject(bucket_name, unique_filename)
        except (BotoCoreError, ClientError) as e:
            log.error(f'Error uploading file: {e}')
            raise BusiException(BusiExceptionEnum.OSS_FILE_UPLOAD_ERROR)

    def get_object(self, object_id: str):
        return self.get_object_from_bucket(self.default_bucket, object_id)

    def get_object_from_bucket(self, bucket_name: str, object_id: str) -> io.IOBase:
        """
        从 AWS S3 中获取对象。

        参数:
        - bucket_name: bucket名称
        - object_id: 对象 ID

        返回:
        - io.IOBase 对象，表示获取的文件。
        """
        try:
            response = self.s3_client.get_object(Bucket=bucket_name, Key=object_id)
            file_stream = io.BytesIO(response['Body'].read())
            file_stream.seek(0)
            log.info(f'File {object_id} successfully retrieved from {bucket_name}')
            return file_stream
        except ClientError as e:
            log.error(f'Error retrieving file: {e}')
            raise BusiException(BusiExceptionEnum.OSS_FILE_RETRIEVAL_ERROR)

    def bucket_exists(self, bucket_name: str) -> bool:
        """
        检查存储桶是否存在。

        参数:
        - bucket_name: 存储桶名称

        返回:
        - 如果存储桶存在，返回 True；否则返回 False。
        """
        try:
            self.s3_client.head_bucket(Bucket=bucket_name)
            return True
        except ClientError as e:
            # 如果存储桶不存在，会抛出 404 错误
            error_code = int(e.response['Error']['Code'])
            if error_code == 404:
                return False
            else:
                raise

    def create_bucket(self, bucket_name: str):
        """
        创建存储桶。

        参数:
        - bucket_name: 存储桶名称
        """
        try:
            self.s3_client.create_bucket(Bucket=bucket_name)
            log.info(f'Bucket {bucket_name} successfully created')
        except (BotoCoreError, ClientError) as e:
            log.error(f'Error creating bucket: {e}')
            raise BusiException(BusiExceptionEnum.OSS_BUCKET_CREATION_ERROR)